package auth

import (
	"errors"
	"fmt"
	"time"

	"github.com/golang-jwt/jwt/v5"
	"github.com/google/uuid"
	"github.com/spf13/viper"
)

type JWTConfig struct {
	AccessSecret  string `yaml:"access_secret" mapstructure:"access_secret"`
	RefreshSecret string `yaml:"refresh_secret" mapstructure:"refresh_secret"`
	AccessExpiry  int    `yaml:"access_expiry" mapstructure:"access_expiry"`
	RefreshExpiry int    `yaml:"refresh_expiry" mapstructure:"refresh_expiry"`
}

type TokenPair struct {
	AccessToken  string `json:"access_token"`
	RefreshToken string `json:"refresh_token"`
}

type CustomClaims struct {
	UserID    string `json:"user_id"`
	Token     string `json:"token"`
	IsRefresh bool   `json:"is_refresh"`
	jwt.RegisteredClaims
}

func GenerateTokenPair(userID string) (*TokenPair, error) {

	// TODO 这里的 jwtConfig，要改成全局的
	var jwtConfig JWTConfig
	err := viper.UnmarshalKey("jwt", &jwtConfig)
	if err != nil {
		return nil, errors.New("jwt config not found")
	}

	accessToken, err := generateToken(userID, jwtConfig.AccessSecret, jwtConfig.AccessExpiry, false)
	if err != nil {
		return nil, err
	}
	refreshToken, err := generateToken(userID, jwtConfig.RefreshSecret, jwtConfig.RefreshExpiry, true)
	if err != nil {
		return nil, err
	}

	return &TokenPair{
		AccessToken:  accessToken,
		RefreshToken: refreshToken,
	}, nil
}

func ValidateAccessToken(token string) (*CustomClaims, error) {
	var jwtConfig JWTConfig
	err := viper.UnmarshalKey("jwt", &jwtConfig)
	if err != nil {
		return nil, errors.New("jwt config not found")
	}

	return validateToken(token, jwtConfig.AccessSecret, false)
}

func ValidateRefreshToken(token string) (*CustomClaims, error) {
	var jwtConfig JWTConfig
	err := viper.UnmarshalKey("jwt", &jwtConfig)
	if err != nil {
		return nil, errors.New("jwt config not found")
	}

	return validateToken(token, jwtConfig.RefreshSecret, true)
}

func generateToken(userID string, secret string, expiry int, isRefresh bool) (string, error) {
	tokenID := uuid.New().String()
	claims := CustomClaims{
		UserID:    userID,
		Token:     tokenID,
		IsRefresh: isRefresh,
		RegisteredClaims: jwt.RegisteredClaims{
			ExpiresAt: jwt.NewNumericDate(time.Now().Add(time.Duration(expiry) * time.Second)),
			IssuedAt:  jwt.NewNumericDate(time.Now()),
		},
	}
	accessToken := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
	return accessToken.SignedString([]byte(secret))
}

func validateToken(tokenString string, secret string, isRefresh bool) (*CustomClaims, error) {
	// 解析令牌
	token, err := jwt.ParseWithClaims(tokenString, &CustomClaims{}, func(token *jwt.Token) (interface{}, error) {
		// 验证签名方法
		if _, ok := token.Method.(*jwt.SigningMethodHMAC); !ok {
			return nil, fmt.Errorf("不支持的签名方法")
		}
		return []byte(secret), nil
	})
	if err != nil {
		return nil, err
	}

	// 验证令牌声明
	if claims, ok := token.Claims.(*CustomClaims); ok && token.Valid {
		if claims.IsRefresh != isRefresh {
			return nil, errors.New("令牌类型不匹配")
		}
		return claims, nil
	}

	return nil, errors.New("无效的令牌")
}
